package com.lenovo.spider.component.scheudler;

import cn.hutool.core.thread.ThreadUtil;
import cn.hutool.crypto.SecureUtil;
import com.alibaba.fastjson.JSON;
import com.lenovo.spider.component.TaskDuplicateChecker;
import com.lenovo.spider.component.TaskIterator;
import com.lenovo.spider.component.duplicate.SimpleTaskDuplicateChecker;
import com.lenovo.spider.domain.Request;
import com.lenovo.spider.domain.Task;
import lombok.extern.slf4j.Slf4j;
import redis.clients.jedis.Jedis;
import redis.clients.jedis.JedisPool;
import redis.clients.jedis.Transaction;
import redis.clients.jedis.Tuple;

import java.io.IOException;
import java.util.List;
import java.util.Set;

/**
 * 基于redis的优先级队列任务迭代器
 *
 * @author <a href="mailto:liuhan3@lenovo.com">HanL(liuhan3)</a>
 * on 2018/4/12 11:44
 */
@Slf4j
public class RedisPriorityTaskIterator implements TaskIterator {

    private final JedisPool jedisPool;
    private final TaskDuplicateChecker duplicateChecker;
    private final QueueConfig queueConfig;

    public RedisPriorityTaskIterator(JedisPool jedisPool, QueueConfig queueConfig) {
        this(jedisPool, queueConfig, new SimpleTaskDuplicateChecker());
    }

    public RedisPriorityTaskIterator(JedisPool jedisPool, QueueConfig queueConfig, TaskDuplicateChecker duplicateChecker) {
        this.jedisPool = jedisPool;
        this.queueConfig = queueConfig;
        this.duplicateChecker = duplicateChecker;
    }

    @Override
    public void addTask(Task task) {
        String queueName = task.getSiteConfig().getIdentity();
        Request request = task.getRequest();
        String queueKey = queueConfig.getQueueKey(queueName);
        String itemKey = queueConfig.getItemKey(queueName);
        if (request.isCheckDuplicate() && duplicateChecker.isDuplicate(task)) {
            log.debug("Task with request: {} is Duplicate, Ignored!", request);
            return;
        }
        try (Jedis jedis = jedisPool.getResource()) {
            jedis.zadd(queueKey, request.getPriority(), request.getUrl());
            jedis.hset(itemKey, SecureUtil.sha1(request.getUrl()), JSON.toJSONString(request));
            log.debug("Task was added to Queue: {}", queueName);
        }
    }

    @Override
    public boolean hasNext() {
        return true;
    }

    @Override
    public Task next() {
        String queueName = queueConfig.getAnyQueueChannel();
        String itemKey = queueConfig.getItemKey(queueName);

        try (Jedis jedis = jedisPool.getResource()) {
            do {
                Tuple tuple = poll(jedis, queueName);
                if (tuple == null) {
                    ThreadUtil.sleep(1000);
                    log.debug("Cannot get any url from Queue, wait for 1s");
                    continue;
                }
                String url = tuple.getElement();
                String field = SecureUtil.sha1(url);
                Request req = JSON.parseObject(jedis.hget(itemKey, field), Request.class);
                if (req == null) {
                    req = Request.builder().url(url).priority((int) tuple.getScore()).build();
                }
                log.info("Make a task with request: {}", req);
                return Task.create(req);
            } while (true);
        }
    }

    /**
     * 从redis的优先级队列中取一个元素
     *
     * @param jedis
     * @param queueName
     * @return
     */
    private Tuple poll(Jedis jedis, String queueName) {
        do {
            jedis.watch(queueName);
            Set<Tuple> set = jedis.zrangeWithScores(queueName, 0, 0);
            if (set.isEmpty()) {
                jedis.unwatch();
                return null;
            }
            Tuple tuple = set.iterator().next();
            String val = tuple.getElement();
            try (Transaction t = jedis.multi()) {
                t.zrem(queueName, val);
                List<Object> results = t.exec();
                long num = results.size() == 0 ? 0 : (long) results.get(0);
                if (num > 0) {
                    return tuple;
                }
            } catch (IOException ignore) {
            }
        } while (true);
    }

}
